# 词形还原并统计词频

import collections
import nltk
import re

from nltk.stem.wordnet import WordNetLemmatizer
from nltk.tokenize import word_tokenize


inputfile = r"..\data\microwave_5reviews.txt" # 评论文本
outputfile = r"..\data\microwave_6extract_word.tsv" # 提取词频

# 正则表达式过滤特殊符号用空格符占位，双引号、单引号、句点、逗号
pat_letter = re.compile(r'[^a-zA-Z \']+')
# 还原常见缩写单词
pat_is = re.compile("(it|he|she|that|this|there|here)(\'s)", re.I)
pat_s = re.compile("(?<=[a-zA-Z])\'s")  # 找出字母后面的字母
pat_s2 = re.compile("(?<=s)\'s?")
pat_not = re.compile("(?<=[a-zA-Z])n\'t")  # not的缩写
pat_would = re.compile("(?<=[a-zA-Z])\'d")  # would的缩写
pat_will = re.compile("(?<=[a-zA-Z])\'ll")  # will的缩写
pat_am = re.compile("(?<=[I|i])\'m")  # am的缩写
pat_are = re.compile("(?<=[a-zA-Z])\'re")  # are的缩写
pat_ve = re.compile("(?<=[a-zA-Z])\'ve")  # have的缩写

lmtzr = WordNetLemmatizer()


def replace_abbreviations(text):
    new_text = text
    new_text = pat_letter.sub(' ', text).strip().lower()
    new_text = pat_is.sub(r"\1 is", new_text)
    new_text = pat_s.sub("", new_text)
    new_text = pat_s2.sub("", new_text)
    new_text = pat_not.sub(" not", new_text)
    new_text = pat_would.sub(" would", new_text)
    new_text = pat_will.sub(" will", new_text)
    new_text = pat_am.sub(" am", new_text)
    new_text = pat_are.sub(" are", new_text)
    new_text = pat_ve.sub(" have", new_text)
    new_text = new_text.replace('\'', ' ')
    return new_text


# pos和tag有相似的地方，通过tag获得pos
def get_wordnet_pos(treebank_tag):
    if treebank_tag.startswith('J'):
        return nltk.corpus.wordnet.ADJ
    elif treebank_tag.startswith('V'):
        return nltk.corpus.wordnet.VERB
    elif treebank_tag.startswith('N'):
        return nltk.corpus.wordnet.NOUN
    elif treebank_tag.startswith('R'):
        return nltk.corpus.wordnet.ADV
    else:
        return ''


def merge(words):
    new_words = []
    for word in words:
        if word:
            tag = nltk.pos_tag(word_tokenize(word))  # tag is like [('bigger', 'JJR')]
            pos = get_wordnet_pos(tag[0][1])
            if pos:
                # lemmatize()方法将word单词还原成pos词性的形式
                lemmatized_word = lmtzr.lemmatize(word, pos)
                new_words.append(lemmatized_word)
            else:
                new_words.append(word)
    return new_words


def get_words(file):
    with open(file, encoding='utf-8') as f:
        words_box = []
        # pat = re.compile(r'[^a-zA-Z \']+') # 过滤特殊符号
        for line in f:
            # if re.match(r'[a-zA-Z]*',line):
            #    words_box.extend(line.strip().strip('\'\"\.,').lower().split())
            # words_box.extend(pat.sub(' ', line).strip().lower().split())
            words_box.extend(merge(replace_abbreviations(line).split()))
    return collections.Counter(words_box)  # 返回单词和词频


def append_ext(words):
    new_words = []
    for item in words:
        word, count = item
        tag = nltk.pos_tag(word_tokenize(word))[0][1]  # tag is like [('bigger', 'JJR')]
        new_words.append((word, count, tag))
    return new_words


# 将统计结果写入文件
#def write_to_file(words, file=r"..\data\microwave_6extract_word.txt"):
def write_to_file(words, file=outputfile):
    f = open(file, 'w', encoding='utf-8')
    for item in words:
        for field in item:
            f.write(str(field) + ',')
        f.write('\n')


if __name__ == '__main__':
    print("counting...")
    # words = get_words(r"..\data\microwave_5reviews.txt")
    words = get_words(inputfile)
    print("writing file...")
    write_to_file(append_ext(words.most_common()))
    print("write success!")
